{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 脱离框架，手写最基本的RAG",
   "id": "e7302b0eebdbaaba"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-07-31T13:14:40.143072Z",
     "start_time": "2025-07-31T13:14:38.719183Z"
    }
   },
   "source": [
    "import numpy as np\n",
    "import requests\n",
    "\n",
    "from config.database import db\n",
    "\n",
    "\n",
    "# 计算余弦相似度\n",
    "def cosine_similarity(vec1, vec2):\n",
    "    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))\n",
    "\n",
    "\n",
    "# 获取 pgvector 嵌入\n",
    "def get_ollama_embedding(text: str) -> list[float]:\n",
    "    payload = {\n",
    "        \"model\": \"dengcao/Qwen3-Embedding-0.6B:Q8_0\",\n",
    "        \"prompt\": text\n",
    "    }\n",
    "    response = requests.post(\"http://localhost:11434/api/embeddings\", json=payload)\n",
    "    if response.status_code == 200:\n",
    "        embedding = response.json()['embedding']\n",
    "        return embedding\n",
    "    else:\n",
    "        raise Exception(f\"Failed to get embedding: {response.text}\")\n",
    "\n",
    "\n",
    "import re\n",
    "\n",
    "\n",
    "# 解析 pgvector 返回的向量字符串\n",
    "def parse_embedding(embedding_str: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    将 pgvector 返回的字符串格式的向量转换为 numpy 数组\n",
    "    示例输入: '[ -3.5594618, 2.123456, 0.123456 ]'\n",
    "    输出: numpy.ndarray([ -3.5594618, 2.123456, 0.123456 ])\n",
    "    \"\"\"\n",
    "    # 去除方括号和空格\n",
    "    clean_str = re.sub(r'[\\[\\]\\s]+', '', embedding_str)\n",
    "    # 分割成浮点数字符串列表\n",
    "    float_strs = clean_str.split(',')\n",
    "    # 转换为浮点数列表\n",
    "    float_list = [float(x) for x in float_strs if x]\n",
    "    # 转换为 numpy 数组（便于计算）\n",
    "    return np.array(float_list, dtype=np.float32)\n",
    "\n",
    "\n",
    "# 检索 知识库\n",
    "def retrieve_knowledge(question: str, threshold=0.75):\n",
    "    embedding = get_ollama_embedding(question)\n",
    "    conn = db.getconn()\n",
    "    cur = conn.cursor()\n",
    "\n",
    "    cur.execute(\"SELECT question, answer, embedding FROM chatbot.public.knowledge\")\n",
    "    rows = cur.fetchall()\n",
    "    best_match = None\n",
    "    best_score = -1\n",
    "\n",
    "    for row in rows:\n",
    "        q, a, db_embedding = row\n",
    "        db_embedding = parse_embedding(db_embedding)  # 关键修改：转换为 float 数组\n",
    "        score = cosine_similarity(embedding, db_embedding)\n",
    "        if score > best_score and score >= threshold:\n",
    "            best_score = score\n",
    "            best_match = a  # 返回答案\n",
    "\n",
    "    cur.close()\n",
    "    db.putconn(conn)\n",
    "\n",
    "    return best_match"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前连接池配置:\n",
      " - 最小连接数 (minconn): 10\n",
      " - 最大连接数 (maxconn): 100\n",
      " - 当前空闲连接数: 10\n",
      " - 当前占用连接数: 0\n",
      "----------------------------------------\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T13:16:37.459152Z",
     "start_time": "2025-07-31T13:16:37.155670Z"
    }
   },
   "cell_type": "code",
   "source": "retrieve_knowledge(\"什么是面向对象编程（OOP）？\", 0.6)",
   "id": "255403d4dd8bf537",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'面向对象编程是一种编程范式，基于类和对象的概念。它支持封装、继承和多态等特性，用于组织和管理复杂的代码结构。'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 手写Standard RAG",
   "id": "8525346fcc089352"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T14:48:43.528736Z",
     "start_time": "2025-07-31T14:48:43.393811Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 新增：基于大模型重写和扩展提示词\n",
    "async def rewrite_and_expand_prompt(original_prompt: str, llm) -> str:\n",
    "    \"\"\"\n",
    "    使用大语言模型对提示词进行改写和扩展，以提升检索效果。\n",
    "\n",
    "    参数:\n",
    "    - original_prompt: 用户的原始问题\n",
    "    - llm: 已初始化的大语言模型实例\n",
    "\n",
    "    返回:\n",
    "    - 改写和扩展后的提示词\n",
    "    \"\"\"\n",
    "    rewritten_prompt = await llm.ainvoke(\n",
    "        f\"请帮我改写并扩展以下问题，使其更适合知识库检索：'{original_prompt}'\"\n",
    "    )\n",
    "    return rewritten_prompt.content\n",
    "\n",
    "\n",
    "from typing import List, Tuple\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def retrieve_knowledge_topk(\n",
    "        question: str,\n",
    "        top_k: int = 3,\n",
    "        threshold: float = 0.75\n",
    ") -> List[Tuple[str, float]]:\n",
    "    \"\"\"\n",
    "    检索与问题相似的知识条目，返回前 top_k 个高分数结果\n",
    "\n",
    "    Args:\n",
    "        question: 用户问题\n",
    "        top_k: 返回的结果数量（默认3个）\n",
    "        threshold: 相似度阈值（仅返回分数>=此值的结果）\n",
    "\n",
    "    Returns:\n",
    "        List[Tuple[answer, score]]: 答案和相似度得分的列表，按分数降序排列\n",
    "    \"\"\"\n",
    "    embedding = get_ollama_embedding(question)\n",
    "    conn = db.getconn()\n",
    "    cur = conn.cursor()\n",
    "\n",
    "    cur.execute(\"SELECT question, answer, embedding FROM chatbot.public.knowledge\")\n",
    "    rows = cur.fetchall()\n",
    "    matches = []  # 存储所有符合条件的候选结果\n",
    "\n",
    "    for row in rows:\n",
    "        q, a, db_embedding = row\n",
    "        db_embedding = parse_embedding(db_embedding)\n",
    "        score = cosine_similarity(embedding, db_embedding)\n",
    "        if score >= threshold:\n",
    "            matches.append((a, score))  # 保存答案和分数\n",
    "\n",
    "    cur.close()\n",
    "    db.putconn(conn)\n",
    "\n",
    "    # 按分数降序排序，并返回前 top_k 个结果\n",
    "    matches.sort(key=lambda x: x[1], reverse=True)\n",
    "    return matches[:top_k]"
   ],
   "id": "736f71bc6abc896c",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T17:44:20.402175Z",
     "start_time": "2025-07-31T17:44:19.880159Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# 加载环境变量\n",
    "load_dotenv()\n",
    "\n",
    "# 初始化模型\n",
    "llm = ChatOpenAI(\n",
    "    api_key=os.getenv(\"DASHSCOPE_KEY\"),\n",
    "    base_url=os.getenv(\"DASHSCOPE_OPENAI_URL\"),\n",
    "    model=\"qwen-plus-1220\"\n",
    ")\n",
    "\n",
    "\n",
    "async def rewrite_to_faq(query: str, llm) -> str:\n",
    "    \"\"\"\n",
    "    将具体查询归纳为常见问题知识库的标准问题\n",
    "    处理规则：\n",
    "    1. 含航班号/日期的具体查询 → 归纳为通用流程类问题\n",
    "    2. 模糊查询（如\"行李规定\"）→ 补充缺失的关键词\n",
    "    3. 禁止返回原始具体信息（如航班号、日期）\n",
    "    \"\"\"\n",
    "    prompt = f\"\"\"\n",
    "    请将以下航空客服查询转换为常见问题知识库的标准问题：\n",
    "    ---\n",
    "    原始查询：{query}\n",
    "    ---\n",
    "    转换规则：\n",
    "    - 若含航班号/日期等具体信息（如CZ3108、3月15日）：\n",
    "      转换为对应的通用流程问题（如\"航班改签具体流程\"）\n",
    "    - 若为模糊查询：\n",
    "      补充缺失的关键词（如\"行李\"→\"托运行李规定\"）\n",
    "    - 禁止包含：\n",
    "      1. 具体航班号/日期/个人信息\n",
    "      2. 解释性文字或示例\n",
    "      3. 禁止生成任何解释、示例、标签或其他非查询内容\n",
    "      4. 输出必须仅为优化后的查询语句，不能包含其他文字\n",
    "    示例：\n",
    "    输入：如何改签3月15日CZ3108航班 → 输出：航班改签具体流程\n",
    "    输入：国际航班能带多少烟 → 输出：国际航班烟草携带规定\n",
    "    \"\"\"\n",
    "    response = await llm.ainvoke(prompt)\n",
    "    return response.content.strip()\n",
    "\n",
    "\n",
    "# 测试示例\n",
    "await rewrite_to_faq(\"行李丢了？\", llm)  # 预期输出：\"托运行李遗失处理流程\""
   ],
   "id": "2034cece7461d113",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'行李丢失处理流程'"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 48
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T17:44:23.354944Z",
     "start_time": "2025-07-31T17:44:22.953120Z"
    }
   },
   "cell_type": "code",
   "source": "await rewrite_to_faq(\"CA123航班延误\", llm)  # 预期输出：\"CA123航班延误\"",
   "id": "4ec230281a688536",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'航班延误如何处理'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 49
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T17:44:25.952854Z",
     "start_time": "2025-07-31T17:44:25.466378Z"
    }
   },
   "cell_type": "code",
   "source": "await rewrite_to_faq(\"国际航班晚点能赔吗\", llm)  # 预期输出：\"国际航班晚点赔偿规定\"",
   "id": "17b37db61fec6073",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'国际航班晚点赔偿规定'"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 50
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-31T17:44:28.942734Z",
     "start_time": "2025-07-31T17:44:28.508622Z"
    }
   },
   "cell_type": "code",
   "source": "await rewrite_to_faq(\"如何改签3月15日CZ3108航班\", llm)",
   "id": "19210b969929a0ac",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'航班改签具体流程'"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 51
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "a69ef32d2ddd35ea"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
